import torch
from tqdm import trange
import numpy as np
from scipy.optimize import linear_sum_assignment as linear_assignment

class K_Means():
    def __init__(self, K: int = 10, device = 'cpu') -> None:
        self.K = K
        self.centres = None
        self.ind = None
        self.device = device

    def centres_initialize(self, x, strategy=None):
        self.centres = torch.rand(self.K, x.shape[1]).to(self.device)  # (K, 784)
        self.centres[self.centres < -1] = -1
        self.centres[self.centres > 1] = 1

    def pairwise_distance(self, x: torch.Tensor, centres: torch.Tensor, strategy=None):
        x = x.unsqueeze(1)
        centres = centres.unsqueeze(0)

        dist = (x - centres) ** 2
        dist = torch.sum(dist, dim=-1).squeeze(-1)
        return dist

    @torch.no_grad()
    def cluster(self, x: torch.Tensor, iter_num):
        '''
        + x (batch, dim): input
        '''
        self.centres_initialize(x)
        
        for it in trange(iter_num):
            dist = self.pairwise_distance(x, self.centres)
            self.ind = torch.argmin(dist, dim=-1)

            for i in range(self.K):
                selected_ind = torch.nonzero(self.ind == i).squeeze()
                selected_x = torch.index_select(x, 0, selected_ind)
                if selected_x.shape[0] == 0:
                    continue
                self.centres[i] = selected_x.mean(dim=0)

    @torch.no_grad()
    def SSE(self, x):
        sum = 0
        for i in range(self.K):
            mask = torch.isin(self.ind, torch.Tensor([i]))
            mask_x = x[mask]
            sum += torch.sum(torch.norm(mask_x - self.centres[i], p=2))

        return sum

    def ACC(self, labels: torch.Tensor):
        y_true = labels.numpy()
        y_pred = self.ind.cpu().numpy()

        y_true = y_true.astype(np.int64)
        assert y_pred.size == y_true.size
        D = max(y_pred.max(), y_true.max()) + 1
        w = np.zeros((D, D), dtype=np.int64)
        for i in range(y_pred.size):
            w[y_pred[i], y_true[i]] += 1
        ind = linear_assignment(w.max() - w)
        ind = np.array(ind).T
        return sum([w[i, j] for i, j in ind]) * 1.0 / y_pred.size

    def info(self, x):
        # print('Clustering Centres[0]:\n', self.centres[0])
        print(self.SSE(x))


def ACC(ind, labels: torch.Tensor):
        y_true = labels.numpy()
        y_pred = ind.cpu().numpy()

        y_true = y_true.astype(np.int64)
        assert y_pred.size == y_true.size
        D = max(y_pred.max(), y_true.max()) + 1
        w = np.zeros((D, D), dtype=np.int64)
        for i in range(y_pred.size):
            w[y_pred[i], y_true[i]] += 1
        ind = linear_assignment(w.max() - w)
        ind = np.array(ind).T
        return sum([w[i, j] for i, j in ind]) * 1.0 / y_pred.size